﻿// -----------------------------------------------------------------------
// <copyright file="PngReader.cs">
// Copyright (c) 2013 Andrey Veselov. All rights reserved.
// License:  Microsoft Public License (MS-PL)
// Contacts: http://andrey.moveax.com  andrey@moveax.com
// </copyright>
// -----------------------------------------------------------------------

namespace PortableImagingLibrary.Png
{
    using System;
    using System.Diagnostics.Contracts;
    using System.IO;
    using System.Linq;
    using System.Threading.Tasks;
    using PortableImagingLibrary.AsyncSupport;
    using PortableImagingLibrary.Convertors;
    using PortableImagingLibrary.Exceptions;
    using PortableImagingLibrary.Png.Compressors;
    using PortableImagingLibrary.Png.Crc;
    using PortableImagingLibrary.Png.DataStructures;
    using PortableImagingLibrary.Png.Filters;

    internal class PngReader
    {
        private readonly byte[] _signature = new byte[] { 137, 80, 78, 71, 13, 10, 26, 10 };
        private const int _headerChunkDataSizeInBytes = 13;
        private const string _headerChunkTypeName = "IHDR";

        private Image _image;
        private PngHeader _header;
        private uint[] _palette;
        private bool _enableCrcCheck;

        ///  <summary>Loads image from the input stream.</summary>
        /// <param name="image">Target image.</param>
        /// <param name="source">The input stream.</param>
        /// <param name="options">The load options.</param>
        /// <exception cref="PortableImagingLibrary.Exceptions.InvalidImageHeaderException">Invalid PNG signature.</exception>
        public async Task<Image> LoadAsync(Image image, Stream source, PngLoadOptions options)
        {
            this._enableCrcCheck = options.EnableCrcCheck;

            using (var reader = new BinaryReaderAsync(source)) {
                // Check file signature 
                byte[] signature = await reader.ReadBytesAsync(8).ConfigureAwait(false);

                if (!this._signature.SequenceEqual(signature))
                    throw new InvalidImageHeaderException("Invalid PNG signature");

                ICrc crc = CrcFactory.Create();

                // Read first chunk (header)
                PngChunk chunk = await this.ReadNextChunkAsync(reader, crc).ConfigureAwait(false);
                this.ProcessHeaderChunk(chunk);
                this._image = image;
                this._image.Reset(this._header.Width, this._header.Height);

                // Create decompressor and filter for data chunks
                using (var decompressor = CompressorFactory.CreateDecompressor(this._header.CompressionType)) {
                    // calc bytes per line
                    const int bitsInByte = 8;
                    int elementsPerSample = this._header.ColorType.ElementsPerSample();

                    decompressor.BytesPerLine = this._header.BitDepth >= 8 ?
                        1 + (this._header.Width * this._header.BitDepth * elementsPerSample) / bitsInByte :
                        1 + (int)Math.Floor(0.5f + this._header.Width * this._header.BitDepth * elementsPerSample / (float)bitsInByte); // +1 per line for control code

                    IFilter filter = FilterFactory.CreateFilter(this._header.FilterType);

                    // Process PNG chunks
                    while (!reader.EndOfStream) {
                        chunk = await this.ReadNextChunkAsync(reader, crc).ConfigureAwait(false);

                        if (string.Equals(chunk.TypeName, "IDAT", StringComparison.Ordinal))
                            this.ProcessImageChunk(chunk, filter, decompressor);
                        else if (string.Equals(chunk.TypeName, "PLTE", StringComparison.Ordinal))
                            this.ProcessPaletteChunk(chunk);
                        else if (string.Equals(chunk.TypeName, "bKGD", StringComparison.Ordinal))
                            this.ProcessBackgroundChunk(chunk);
                        else if (string.Equals(chunk.TypeName, "tRNS", StringComparison.Ordinal))
                            this.ProcessTransparencyChunk(chunk);
                    }
                }
            }

            return this._image;
        }

        #region Chunk reader methods

        private async Task<PngChunk> ReadNextChunkAsync(BinaryReaderAsync reader, ICrc crc)
        {
            var chunk = new PngChunk {
                DataSize = await reader.ReadIntAsync().ConfigureAwait(false),
                TypeName = await reader.ReadStringAsync(4).ConfigureAwait(false)
            };

            chunk.Data = await reader.ReadBytesAsync(chunk.DataSize).ConfigureAwait(false);
            chunk.Crc = await reader.ReadUintAsync().ConfigureAwait(false);

            // Validate CRC
            if (this._enableCrcCheck) {
                crc.Reset();

                crc.Update(chunk.TypeName);
                crc.Update(chunk.Data);

                if (chunk.Crc != crc.CurrentValue)
                    throw new InvalidImageFormatException(string.Format("Invalid chunk ({0}) CRC", chunk.TypeName));
            }

            return chunk;
        }

        private void ProcessHeaderChunk(PngChunk chunk)
        {
            if (!string.Equals(chunk.TypeName, PngReader._headerChunkTypeName) || chunk.DataSize != PngReader._headerChunkDataSizeInBytes)
                throw new InvalidImageHeaderException("Header data was not found or has invalid length.");

            this._header = new PngHeader {
                Width = ByteConvertor.BytesToInt(chunk.Data, 0),
                Height = ByteConvertor.BytesToInt(chunk.Data, 4),
                BitDepth = chunk.Data[8],
                ColorType = (ColorType)chunk.Data[9],
                CompressionType = (CompressionType)chunk.Data[10],
                FilterType = (FilterType)chunk.Data[11],
                InterlaceType = (InterlaceType)chunk.Data[12]
            };

            // Validate PNG properties for current support state
            if (this._header.BitDepth > 8)
                throw new NotSupportedException("Images with BitDepth more then 8 are not supported at this time.");

            if (this._header.ColorType == ColorType.Greyscale || this._header.ColorType == ColorType.GreyscaleWithAlpha)
                throw new NotSupportedException("Grayscale images are not supported at this time.");

            if (this._header.InterlaceType == InterlaceType.Adam7)
                throw new NotSupportedException("Interlaced images are not supported at this time.");
        }

        private void ProcessPaletteChunk(PngChunk chunk)
        {
            Contract.Requires(chunk != null);
            Contract.Requires(chunk.Data != null);
            Contract.Requires(chunk.Data.Length >= 3);

            int count = chunk.Data.Length / 3;

            this._palette = new uint[count];
            for (int index = 0; index < count; index++)
                this._palette[index] = ColorConvertor.RgbBytesTo8BitColor(chunk.Data, index * 3);
        }

        private void ProcessBackgroundChunk(PngChunk chunk)
        {
            switch (this._header.ColorType) {
                case ColorType.Greyscale:
                case ColorType.GreyscaleWithAlpha:
                    throw new NotSupportedException();

                case ColorType.Truecolor:
                case ColorType.TruecolorWithAlpha:
                    this._image.BackgroundColor = ColorConvertor.Conver16BitArrayTo8BitColor(chunk.Data, 0, this._header.BitDepth <= 8);
                    break;

                case ColorType.IndexedColor:
                    this._image.BackgroundColor = this._palette[chunk.Data[0]];
                    break;

                default:
                    throw new InvalidImageFormatException("invalid colorType");
            }
        }

        private void ProcessTransparencyChunk(PngChunk chunk)
        {
            switch (this._header.ColorType) {
                case ColorType.Truecolor:
                    throw new NotImplementedException();

                case ColorType.IndexedColor:
                    for (var index = 0; index < chunk.Data.Length; index++)
                        this._palette[index] = (uint)ByteConvertor.ReplaceByte((int)this._palette[index], 3, chunk.Data[index]);
                    break;

                case ColorType.Greyscale:
                    throw new NotSupportedException();

                default:
                    throw new InvalidImageFormatException("Transparency chunk should be present only for truecolor, greyscale or indexed images.");
            }
        }

        private void ProcessImageChunk(PngChunk chunk, IFilter filter, IDecompressor decompressor)
        {
            decompressor.SetCurrentChunkData(chunk.Data);
            var pixelReader = this.GetPixelReader();
            int y = 0;
            var pixels = this._image.Pixels;

            while (decompressor.ReadLine()) {

                filter.ReconstructLine(decompressor.OutputBuffer, (byte)this._header.ColorType.ElementsPerSample());

                for (var x = 0; x < this._header.Width; x++)
                    pixels[x + y * this._header.Width] = pixelReader(x, decompressor.OutputBuffer);

                y++;
            }
        }

        #endregion Chunk reader methods

        #region Pixel readers

        private Func<int, byte[], uint> GetPixelReader()
        {
            switch (this._header.ColorType) {
                case ColorType.Truecolor:
                    return this.PixelReaderRgb;

                case ColorType.TruecolorWithAlpha:
                    return this.PixelReaderArgb;

                case ColorType.IndexedColor:
                    return this.PixelReaderIndexed;

                default:
                    throw new NotSupportedException();
            }
        }

        private uint PixelReaderArgb(int x, byte[] lineBuffer)
        {
            int xBase = 1 + x * 4;
            return (uint)(lineBuffer[xBase + 3] << 24 | lineBuffer[xBase] << 16 | lineBuffer[xBase + 1] << 8 | lineBuffer[xBase + 2]);
        }

        private uint PixelReaderRgb(int x, byte[] lineBuffer)
        {
            int xBase = 1 + x * 3;
            return (uint)(255 << 24 | lineBuffer[xBase] << 16 | lineBuffer[xBase + 1] << 8 | lineBuffer[xBase + 2]);
        }

        private uint PixelReaderIndexed(int x, byte[] lineBuffer)
        {
            if (this._header.BitDepth == 8)
                return this._palette[lineBuffer[x + 1]];

            int bitPos = 8 + x * this._header.BitDepth; // add 8 to skip control code in line buffer
            int bytePos = bitPos / 8;
            int bitInBytePos = bitPos % 8;

            byte paletteIndex = lineBuffer[bytePos];

            paletteIndex >>= (8 - (bitInBytePos + this._header.BitDepth));
            paletteIndex <<= (8 - this._header.BitDepth);
            paletteIndex >>= (8 - this._header.BitDepth);

            return this._palette[paletteIndex];
        }

        #endregion Pixel readers
    }
}
